/*
 * dhmm.cc
 *
 *  Created on: Aug 11, 2011
 *      Author: tqlong
 */

#include <algorithm>
#include "dhmm.h"

#define ZERO log(0)
#define ONE log(1)

struct LogSpace
{
  typedef arma::vec dvec;

  static double
  mul(double a, double b)
  {
    return a + b;
  }

  static double
  div(double a, double b)
  {
    return a - b;
  }

  static double
  add(double a, double b)
  {
    if (a == -std::numeric_limits<double>::infinity())
      return b;
    else if (b == -std::numeric_limits<double>::infinity())
      return a;
    else if (a > b)
      return a + log1p(exp(b - a));
    else
      return b + log1p(exp(a - b));
  }

  static double
  vecsum(dvec& v)
  {
    double s = ZERO;
    for (unsigned int i = 0; i < v.n_elem; i++)
      {
        s = add(s, v[i]);
        //std::cout << "+s=" << s << std::endl;
      }
    return s;
  }

  static void
  vecscale(dvec& v, double s)
  {
    for (unsigned int i = 0; i < v.n_elem; i++)
      v(i) = mul(v(i), s);
  }

  static void
  vecnormalize(dvec& v)
  {
    //v.print("vp=");
    double s = vecsum(v);
    //std::cout << "s=" << s << std::endl;
    vecscale(v, -s);
    //v.print("v =");
  }

  /** Pick a random position in a vector of probability, avoid position with zero probability */
  template<typename V>
    static int
    randomWheel(const V& v)
    {
      double s = ZERO;

      for (unsigned int i = 0; i < v.n_elem; i++)
        s = add(s, v[i]);
      //std::cout << "s=" << s << std::endl;
      if (s == std::numeric_limits<double>::infinity())
        return -1; // no position is picked

      // pick a random number and track the cummulative sum
      double r = mul(log(arma::randu()), s), cs = ZERO;
      //std::cout << "r=" << r << std::endl;
      for (unsigned int i = 0; i < v.n_elem; i++)
        {
          cs = add(cs, v[i]);
          if (r <= cs)
            return i;
        }

      return -1;
    }

};

DurationHMM::DurationHMM(int n_state, int n_observation) :
  Q_(n_state), K_(n_observation)
{
  setTransition(arma::randu<dmat>(Q_, Q_));
  setEmission(arma::randu<dmat>(K_, Q_));
  setStart(arma::randu<dvec>(Q_));
  d_ = std::vector<ivec>(Q_);
  p_ = std::vector<dvec>(Q_);
}

void
DurationHMM::setTransition(const dmat& aa)
{
  a_ = arma::log(aa);
  for (unsigned int i = 0; i < a_.n_cols; i++)
    {
      dvec v = a_.col(i);
      LogSpace::vecnormalize(v);
      a_.col(i) = v;
    }
}

void
DurationHMM::setEmission(const dmat& bb)
{
  b_ = arma::log(bb);
  for (unsigned int i = 0; i < b_.n_cols; i++)
    {
      dvec v = b_.col(i);
      LogSpace::vecnormalize(v);
      b_.col(i) = v;
    }
}

void
DurationHMM::setStart(const dvec& pi)
{
  pi_ = arma::log(pi);
  LogSpace::vecnormalize(pi_);
}

void
DurationHMM::setDuration(int state, const ivec& duration, const dvec& prob)
{
  d_[state] = duration;
  p_[state] = prob;
  LogSpace::vecnormalize(p_[state]);
}

void
DurationHMM::dump()
{
  a_.print("a = ");
  b_.print("b = ");
  pi_.print("pi = ");
  for (int i = 0; i < Q_; i++)
    {
      std::cout << "state i" << std::endl;
      d_[i].print("d = ");
      p_[i].print("p = ");
    }
}

void
DurationHMM::genSeq(int T, ivec& o, ivec& q, ivec& d)
{
  d = q = ivec(T);
  q[0] = LogSpace::randomWheel(pi_);
  //std::cout << "q[0] = " << q[0] << std::endl;
  d[0] = d_[q[0]][LogSpace::randomWheel(p_[q[0]])];
  //std::cout << "d[0] = " << d[0] << std::endl;
  int L = d[0];
  for (int t = 1; t < T; t++)
    {
      q[t] = LogSpace::randomWheel(a_.col(q[t - 1]));
      d[t] = d_[q[t]][LogSpace::randomWheel(p_[q[t]])];
      L += d[t];
    }
  o = ivec(L);
  int l = 0;
  for (int t = 0; t < T; t++)
    {
      for (int qd = 0; qd < d[t]; qd++, l++)
        {
          o[l] = LogSpace::randomWheel(b_.col(q[t]));
        }
    }
}

double
DurationHMM::seqProbability(const ivec& o)
{
  computeAlpha(o);

  int T = o.n_elem;
  double s = ZERO;
  for (int m = 0; m < Q_; m++)
    {
      s = LogSpace::add(s, alpha_[m][T - 1]);
    }

  //  computeBeta(o);
  //  double ss = ZERO;
  //  for (int m = 0; m < Q_; m++)
  //    for (unsigned int d_idx = 0; d_idx < d_[m].n_elem; d_idx++)
  //      {
  //        int d = d_[m][d_idx];
  //        if (d > T) continue;
  //        double tmp = LogSpace::mul(pi_[m], p_[m][d_idx]);
  //        tmp = LogSpace::mul(tmp, dprob(o, d-1, m, d));
  //        tmp = LogSpace::mul(tmp, beta_[m][d-1]);
  //        ss = LogSpace::add(ss, tmp);
  //      }
  //  std::cout << "ss = " << ss << std::endl;

  return s;
}

double
DurationHMM::dprob(const ivec& o, int t, int m, int d)
{
  double s = ONE;
  for (int i = t - d + 1; i <= t; i++)
    {
      s = LogSpace::mul(s, b_(o[i], m));
    }
  return s;
}

void
DurationHMM::computeAlpha(const ivec& o)
{
  int T = o.n_elem;

  // allocate alpha
  alpha_ = dvecvec(Q_);
  for (int m = 0; m < Q_; m++)
    alpha_[m] = dvec(T);

  for (int t = 0; t < T; t++)
    {
      //std::cout << "--t = " << t << std::endl;
      for (int m = 0; m < Q_; m++)
        {
          //std::cout << "----m = " << m << std::endl;
          double s = ZERO;
          for (unsigned int d_idx = 0; d_idx < d_[m].n_elem; d_idx++)
            {
              int d = d_[m][d_idx];
              if (t - d + 1 < 0)
                continue;
              double tmp = ZERO;
              if (t - d + 1 == 0)
                {
                  tmp = LogSpace::mul(pi_[m], p_[m][d_idx]);
                  tmp = LogSpace::mul(tmp, dprob(o, t, m, d));
                }
              else
                {
                  for (int n = 0; n < Q_; n++)
                    if (n != m)
                      {
                        double tmp1 = LogSpace::mul(alpha_[n][t - d], a_(m, n));
                        tmp1 = LogSpace::mul(tmp1, p_[m][d_idx]);
                        tmp1 = LogSpace::mul(tmp1, dprob(o, t, m, d));
                        tmp = LogSpace::add(tmp, tmp1);
                      }
                }
              s = LogSpace::add(s, tmp);
            }
          alpha_[m][t] = s;
        }
    }
}

void
DurationHMM::computeBeta(const ivec& o)
{
  int T = o.n_elem;

  // allocate alpha
  beta_ = dvecvec(Q_);
  for (int m = 0; m < Q_; m++)
    beta_[m] = dvec(T);

  for (int t = T - 1; t >= 0; t--)
    {
      for (int m = 0; m < Q_; m++)
        {
          if (t == T - 1)
            {
              beta_[m][t] = ONE;
              continue;
            }
          double s = ZERO;
          for (int n = 0; n < Q_; n++)
            if (n != m)
              for (unsigned int dn_idx = 0; dn_idx < d_[n].n_elem; dn_idx++)
                {
                  int dn = d_[n][dn_idx];
                  if (t + dn >= T)
                    continue;
                  double tmp = LogSpace::mul(a_(n, m), p_[n][dn_idx]);
                  tmp = LogSpace::mul(tmp, dprob(o, t + dn, n, dn));
                  tmp = LogSpace::mul(tmp, beta_[n][t + dn]);
                  s = LogSpace::add(s, tmp);
                }
          beta_[m][t] = s;
        }
    }
}

void
DurationHMM::computeAlphaS(const ivec& o)
{
  int T = o.n_elem;

  // allocate alpha
  alphaS_ = dvecvec(Q_);
  for (int m = 0; m < Q_; m++)
    alphaS_[m] = dvec(T);

  for (int t = 0; t < T; t++)
    {
      for (int m = 0; m < Q_; m++)
        {
          double s = ZERO;
          for (int n = 0; n < Q_; n++)
            if (n != m)
              {
                double tmp = LogSpace::mul(alpha_[n][t], a_(m, n));
                s = LogSpace::add(s, tmp);
              }
          alphaS_[m][t] = s;
        }
    }
}

void
DurationHMM::computeBetaS(const ivec& o)
{
  int T = o.n_elem;

  // allocate alpha
  betaS_ = dvecvec(Q_);
  for (int m = 0; m < Q_; m++)
    betaS_[m] = dvec(T);

  for (int t = 0; t < T; t++)
    {
      for (int m = 0; m < Q_; m++)
        {
          double s = ZERO;
          for (unsigned int d_idx = 0; d_idx < d_[m].n_elem; d_idx++)
            {
              int d = d_[m][d_idx];
              if (t + d >= T)
                continue;
              double tmp = LogSpace::mul(beta_[m][t + d], p_[m][d_idx]);
              tmp = LogSpace::mul(tmp, dprob(o, t + d, m, d));
              s = LogSpace::add(s, tmp);
            }
          betaS_[m][t] = s;
        }
    }

  beta0S_ = dvec(Q_);
  for (int m = 0; m < Q_; m++)
    {
      double s = ZERO;
      for (unsigned int d_idx = 0; d_idx < d_[m].n_elem; d_idx++)
        {
          int d = d_[m][d_idx];
          if (d > T)
            continue;
          double tmp = LogSpace::mul(beta_[m][d - 1], p_[m][d_idx]);
          tmp = LogSpace::mul(tmp, dprob(o, d - 1, m, d));
          s = LogSpace::add(s, tmp);
        }
      beta0S_[m] = s;
    }
}

void
DurationHMM::train(const ivec& o)
{
  for (int iter = 0; iter < 1; iter++)
    {
      computeAlpha(o);
      computeAlphaS(o);
      computeBeta(o);
      computeBetaS(o);

      int T = o.n_elem;
      double s = ZERO;
      for (int m = 0; m < Q_; m++)
        {
          s = LogSpace::add(s, alpha_[m][T - 1]);
        }
      std::cout << "s = " << s << std::endl;

      reestimateP(o);
      reestimateB(o);
    }
  std::cout << "s = " << seqProbability(o) << std::endl;
}

void
DurationHMM::reestimateP(const ivec& o)
{
  dvecvec acc_pd(Q_);
  dvec acc_p(Q_);
  acc_p.fill(ZERO);
  for (int m = 0; m < Q_; m++)
    {
      acc_pd[m] = p_[m];
      acc_pd[m].fill(ZERO);
    }

  int T = o.n_elem;
  for (int t = 0; t < T; t++)
    for (int m = 0; m < Q_; m++)
      for (unsigned int d_idx = 0; d_idx < d_[m].n_elem; d_idx++)
        {
          int d = d_[m][d_idx];
          if (t + d >= T)
            continue;
          double tmp = LogSpace::mul(alphaS_[m][t], p_[m][d_idx]);
          tmp = LogSpace::mul(tmp, beta_[m][t + d]);
          tmp = LogSpace::mul(tmp, dprob(o, t + d, m, d));
          acc_pd[m][d_idx] = LogSpace::add(acc_pd[m][d_idx], tmp);
          acc_p[m] = LogSpace::add(acc_p[m], tmp);
        }

  for (int m = 0; m < Q_; m++)
    for (unsigned int d_idx = 0; d_idx < d_[m].n_elem; d_idx++)
      {
        p_[m][d_idx] = LogSpace::div(acc_pd[m][d_idx], acc_p[m]);
      }
}

void
DurationHMM::reestimateB(const ivec& o)
{
  int T = o.n_elem;
  dmat acc_bk(K_, Q_);
  dvec acc_b(Q_);
  dvec ab(T), abS(T);
  acc_bk.fill(ZERO);
  acc_b.fill(ZERO);
  ab.fill(ZERO);
  abS.fill(ZERO);

  for (int t = 0; t < T; t++)
    {
      int k = o[t];
      for (int m = 0; m < Q_; m++)
        {

        }
    }
}
